import mindspore as ms
import mindspore.dataset as ds
from util.datasets import synthetic_data
from mindspore.dataset.vision import HWC2CHW
from mindspore.dataset.transforms import TypeCast


def data_iter(features, labels, batch_size):
    data_gen = synthetic_data(features, labels)
    data = ds.GeneratorDataset(tuple(data_gen), ['image', 'label'], shuffle=True)
    data = data.map(operations=[TypeCast(ms.dtype.float32)])
    data = data.map(operations=[TypeCast(ms.dtype.int32)],
                    input_columns='label')
    data = data.batch(batch_size)
    return data


def data_iter_2d(features, labels, batch_size):
    data_gen = synthetic_data(features, labels)
    data = ds.GeneratorDataset(tuple(data_gen), ['image', 'label'], shuffle=True)
    data = data.map(operations=[TypeCast(ms.dtype.float32),
                                HWC2CHW()])
    data = data.map(operations=[TypeCast(ms.dtype.int32)],
                    input_columns='label')
    data = data.batch(batch_size)
    return data
